package main

import (
	"bufio"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/google/gopacket/pfring"
	"net"
	"os"
	"strconv"
	"time"
)

/**
  基于pfring+gopacket 抓包解包发包
  单个IP阻断运行命令 ./libpcapAccess.go [pfring集群节点] [要阻断的IP] [设置BPF]
  概念: 集群,跨多个环进行平衡时使用的一种集群类型

*/
var (
	eth   layers.Ethernet
	dot1q layers.Dot1Q
	ipv4  layers.IPv4
	tcp   layers.TCP
	ringR *pfring.Ring
	//ringW *pfring.Ring
	err  error
	ip   net.IP
	node int
)

func main() {
	// 获取命令行参数
	//1.0 获取第二个参数,设置节点
	node, err = strconv.Atoi(os.Args[1])
	if err != nil {
		panic("please enter the correct number of nodes,not " + strconv.Itoa(node))
	}
	//2.0 获取第三个参数,设置阻断IP,单个
	ip = net.ParseIP(os.Args[2])
	if ip == nil {
		panic("please enter the correct ip not " + ip.String())
	}
	//3.0 获取第四个参数,设置BPF
	bpf := os.Args[3]

	//4.0 读取控制台输入的网卡
	devs, _ := pcap.FindAllDevs()
	fmt.Println("Please enter the network card, for example eth0")
	for i := range devs {
		fmt.Println(devs[i].Name)
	}
	fmt.Println("Please enter :")
	in := bufio.NewReader(os.Stdin)
	str := make([]byte, 1024)
	str, _, err = in.ReadLine()
	if err != nil {
		panic(err)
	}
	deviceId := string(str)
	//设置读取数据包的ring
	ringR, err = pfring.NewRing(deviceId, 65535, pfring.FlagPromisc)
	if err != nil {
		panic(err)
	}
	if bpf != "" {
		err = ringR.SetBPFFilter(bpf)
	} else {
		err = ringR.SetBPFFilter("tcp")
	}
	if err != nil { // optional
		panic(err)
	}
	//err = ringR.SetSocketMode(pfring.ReadOnly)
	//if err != nil {
	//	panic(err)
	//	return
	//}
	err = ringR.Enable()
	if err != nil { // Must do this!, or you get no packets!
		panic(err)
	}
	err = ringR.SetCluster(node, pfring.ClusterPerFlow)
	//设置写包的ring
	//ringW, err = pfring.NewRing("enp4s0", 65535, pfring.FlagPromisc)
	//if err != nil {
	//	panic(err)
	//}
	//err = ringR.SetSocketMode(pfring.WriteOnly)
	//if err != nil {
	//	panic(err)
	//	return
	//}
	//err = ringW.Enable()
	//if err != nil { // Must do this!, or you get no packets!
	//	panic(err)
	//}
	fmt.Println("The program has been started")
	fmt.Println("当前节点:", node)
	fmt.Println("当前阻断IP:", ip)
	fmt.Println("当前BPF:", bpf)
	fmt.Println("当前deviceId:", deviceId)

	dlc := gopacket.DecodingLayerContainer(gopacket.DecodingLayerArray(nil))
	dlc = dlc.Put(&eth)
	dlc = dlc.Put(&dot1q)
	dlc = dlc.Put(&ipv4)
	dlc = dlc.Put(&tcp)
	parser := dlc.LayersDecoder(layers.LayerTypeEthernet, gopacket.NilDecodeFeedback)
	decoded := make([]gopacket.LayerType, 0, 20)
	for true {
		data, _, _ := ringR.ReadPacketData()
		analysePacket(data, parser, decoded)
	}
	//packetSource := gopacket.NewPacketSource(ring, layers.LinkTypeEthernet)
	//for packet := range packetSource.Packets() {
	//	analysePacket(packet.Data(), parser, decoded)
	//}
}

func analysePacket(packet []byte, parser gopacket.DecodingLayerFunc, decoded []gopacket.LayerType) {
	lt, err := parser(packet, &decoded)
	if err != nil {
		//fmt.Fprintf(os.Stderr, "无法解码层：%v\n", err)
		return
	}
	if lt != gopacket.LayerTypeZero {
		//fmt.Fprintf(os.Stderr, "未知层类型: %v\n", lt)
		return
	}
	fmt.Println("1srcIp: ", ipv4.SrcIP.String(), "dstIp: ", ipv4.DstIP.String())
	if !ip.Equal(ipv4.SrcIP) && !ip.Equal(ipv4.DstIP) {
		return
	}
	fmt.Println("2srcIp: ", ipv4.SrcIP.String(), "dstIp: ", ipv4.DstIP.String())
	//5.0 过滤掉SYN FIN RST 和 http
	if tcp.SYN || tcp.FIN || tcp.RST {
		//log.Println("send nothing")
		return
	}

	//握手完成开始推送数据时候拦截
	for i := 0; i < 5; i++ {
		err = SendFlag(eth.DstMAC, eth.SrcMAC, ipv4.DstIP, ipv4.SrcIP, tcp.DstPort, tcp.SrcPort, tcp.Ack, 10000, false, false, true)
		if err != nil {
			panic(err)
		}
	}
	//SendFlag(eth.SrcMAC, eth.DstMAC, ipv4.SrcIP, ipv4.DstIP, tcp.SrcPort, tcp.DstPort, tcp.Seq, tcp.Ack, false, false, true)
	fmt.Println("send rst")
}

func SendFlag(srcMac, dstMac net.HardwareAddr, srcIp, dstIp net.IP, srcPort, dstPort layers.TCPPort, seq, ack uint32, SYN, ACK, RST bool) error {
	s := time.Now().UnixNano()
	eth.SrcMAC = srcMac
	eth.DstMAC = dstMac
	ipv4.SrcIP = srcIp
	ipv4.DstIP = dstIp
	tcp.SrcPort = srcPort
	tcp.DstPort = dstPort
	tcp.SYN = SYN
	tcp.Seq = seq
	tcp.ACK = ACK
	tcp.Ack = ack
	tcp.RST = RST

	options := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}

	if err = tcp.SetNetworkLayerForChecksum(&ipv4); err != nil {
		return err
	}
	buffer := gopacket.NewSerializeBuffer()

	err = gopacket.SerializeLayers(buffer, options, &eth, &ipv4, &tcp)
	if err != nil {
		return err
	}
	//resetPacket := gopacket.NewPacket(buffer.Bytes(), layers.LayerTypeEthernet, gopacket.Default)
	err = ringR.WritePacketData(buffer.Bytes())
	if err != nil {
		return err
	}
	fmt.Println("send time", time.Now().UnixNano()-s)
	return nil
}
